#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr

import torch
from scene import Scene
import os
from tqdm import tqdm
from os import makedirs
from gaussian_renderer import render
import torchvision
from utils.general_utils import safe_state
from argparse import ArgumentParser
from hparam import HyperParams, get_combined_args
import time
import json
import numpy as np

def render_set(hp, scene, mode, views, background):
    render_path = os.path.join(scene.model_path, mode, "ours_{}".format(scene.load_iter), "renders")
    gts_path = os.path.join(scene.model_path, mode, "ours_{}".format(scene.load_iter), "gt")

    makedirs(render_path, exist_ok=True)
    makedirs(gts_path, exist_ok=True)

    t_list = []
    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        torch.cuda.synchronize();t0 =time.time()
        freq_images = []
        for freq_idx in scene.all_gaussians.keys():
            gaussians = scene.activate_gaussian(freq_idx)
            freq_images.append(render(gaussians, view, background)['render'])
        rendering = sum(freq_images)
        torch.cuda.synchronize();t1= time.time()
        t_list.append(t1-t0)
        
        gt = view.gt_image[0:3, :, :]
        torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(idx) + ".png"))
        torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(idx) + ".png"))
        
    t = np.array(t_list[5:])
    fps = 1.0 / t.mean()
    print(f'Test FPS: \033[1;35m{fps:.5f}\033[0m')
    
    n_points=[]
    for freq_idx in scene.all_gaussians.keys():
        gaussians = scene.activate_gaussian(freq_idx)
        n_points.append(gaussians.n_points)
        
    if not os.path.exists(hp.model_path + "/results.json"):
        with open(hp.model_path + "/results.json", 'w') as fp:
            json.dump({}, fp)  # 创建一个空字典的 JSON 文件
    with open(hp.model_path + "/results.json", 'r') as file:
        data = json.load(file)
    data.update({
        "Test FPS": fps,
        'Pts': f"{n_points[0]}|{n_points[1]}|{sum(n_points)}"
    })
    with open(hp.model_path + "/results.json", 'w') as file:
        json.dump(data, file, indent=True)
    
def render_sets(args, hp):
    with torch.no_grad():
        scene = Scene(hp, args.load_iter, args.training)
        background = scene.random_background() if hp.random_background else scene.background

        # if not args.skip_train:
        #      render_set(hp, scene, "train", scene.get_train_cameras(), background)

        if not args.skip_test:
             render_set(hp, scene, 'test', scene.get_test_cameras(), background)

if __name__ == "__main__":
    # Set up command line argument parser
    parser = ArgumentParser(description="Testing script parameters")
    parser.add_argument("--load_iter", default=-1, type=int)
    parser.add_argument("--skip_train", action="store_true")
    parser.add_argument("--skip_test", action="store_true")
    parser.add_argument('--detect_anomaly', action='store_true')
    parser.add_argument('--training', action='store_true')
    parser.add_argument("--quiet", action="store_true")
    args, _ = parser.parse_known_args()
    
    # Initialize system state (RNG)
    safe_state(args.quiet)
    torch.autograd.set_detect_anomaly(args.detect_anomaly)
    
    hp = HyperParams()
    
    hp.send_to(parser)
    args = parser.parse_args()
    hp.extract_from(args)
    hp = get_combined_args(args, hp)
    
    print('Rendering:', hp.model_path)
    render_sets(args, hp)